import torch  
import torch.nn as nn  
import torch.optim as optim  
from torch.utils.data import DataLoader  
from torchvision import datasets, transforms  

# 1. 设置超参数  
batch_size = 64       # 批大小  
learning_rate = 0.0001    # 学习率  
num_epochs = 10        # 训练轮次  

# 2. 数据预处理  
transform = transforms.Compose([  
    transforms.ToTensor(),                  # 将图像转换为张量  
    transforms.Normalize((0.5,), (0.5,))    # 对图像进行归一化处理  
])  

# 3. 下载并加载训练数据和测试数据  
train_dataset = datasets.MNIST(root='data', train=True, transform=transform, download=True)  
test_dataset = datasets.MNIST(root='data', train=False, transform=transform, download=True)  
train_loader = DataLoader(dataset=train_dataset, batch_size=batch_size, shuffle=True)  
test_loader = DataLoader(dataset=test_dataset, batch_size=batch_size, shuffle=False)  

# 4. 检查 CUDA 是否可用，并设置设备  
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')  
print(f'Using device: {device}')  

# 4. 定义神经网络模型  
class SimpleNN(nn.Module):  
    def __init__(self):  
        super(SimpleNN, self).__init__()  
        self.fc1 = nn.Linear(28 * 28, 128)  # 输入层，将 28x28 图像展平为 784  
        self.fc2 = nn.Linear(128, 64)        # 隐藏层  
        self.fc3 = nn.Linear(64, 10)         # 输出层，10 类（从 0 到 9）  

    def forward(self, x):  
        x = x.view(-1, 28 * 28)            # 将图像展平  
        x = torch.relu(self.fc1(x))        # 第 1 层 ReLU 激活  
        x = torch.relu(self.fc2(x))        # 第 2 层 ReLU 激活  
        x = self.fc3(x)                     # 输出层，不需要激活函数  
        return x  

# 5. 实例化模型、损失函数和优化器  
model = SimpleNN()                        # 创建模型实例  
criterion = nn.CrossEntropyLoss()        # 交叉熵损失函数  
optimizer = optim.Adam(model.parameters(), lr=learning_rate)  # Adam 优化器  

# 6. 训练模型  
for epoch in range(num_epochs):  
    for images, labels in train_loader:  
        optimizer.zero_grad()             # 清空梯度  
        outputs = model(images)           # 获取模型输出  
        loss = criterion(outputs, labels) # 计算损失  
        loss.backward()                   # 反向传播  
        optimizer.step()                  # 更新参数  

    print(f'Epoch [{epoch+1}/{num_epochs}], Loss: {loss.item():.4f}')  # 打印损失  

# 7. 测试模型  
model.eval()  # 设置模型为评估模式  
with torch.no_grad():  # 不计算梯度  
    correct = 0  
    total = 0  
    for images, labels in test_loader:  
        outputs = model(images)            # 获取模型输出  
        _, predicted = torch.max(outputs.data, 1)  # 获取预测类  
        total += labels.size(0)            # 更新总样本数  
        correct += (predicted == labels).sum().item()  # 更新正确预测的样本数  

    print(f'准确率：{100 * correct / total:.2f}%')  # 打印模型准确率